using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Globalization;
using System.IO;

using Atomic.Libraries;
using Atomic.Libraries.Compression;
using Atomic.Structures;
using Atomic.Vasp;

namespace Atomic.Interfaces.Phonopy
{
	[Serializable]
	public class PhonopyProblem
	{
		private PhonopyProblem(Structure baseStructure, LatticeRepeatMatrix repeatMatrix, double? displacement, Structure superStructure, IEnumerable<Structure> displacedStructures, CompressedData disp, CompressedData poscar)
		{
			BaseStructure = baseStructure;
			RepeatMatrix = repeatMatrix;
			Displacement = displacement;
			SuperStructure = superStructure;
			DisplacedStructures = new ReadOnlyList<Structure>(displacedStructures);
			Disp = disp;
			Poscar = poscar;
		}

		public static PhonopyProblem Generate(Structure baseStructure, LatticeRepeatMatrix repeatMatrix)
		{
			return Generate(baseStructure, repeatMatrix, null);
		}

		public static PhonopyProblem Generate(Structure baseStructure, LatticeRepeatMatrix repeatMatrix, double? displacement)
		{
			using (TempDirectory temp = new TempDirectory())
			{
				return Generate(baseStructure, repeatMatrix, displacement, temp.Directory);
			}
		}

		private static PhonopyProblem Generate(Structure baseStructure, LatticeRepeatMatrix repeatMatrix, double? displacement, DirectoryInfo directory)
		{
			VaspPoscar.Write(baseStructure, directory.AddFile("POSCAR"));

			Process p = new Process();
			p.StartInfo.FileName = "phonopy";
			if (displacement.HasValue)
			{
				p.StartInfo.Arguments = string.Format(CultureInfo.InvariantCulture, "-d --dim='{0}' --amplitude={1:r}", FormatRepeatMatrix(repeatMatrix), displacement.Value);
			}
			else
			{
				p.StartInfo.Arguments = string.Format("-d --dim='{0}'", FormatRepeatMatrix(repeatMatrix));
			}
			p.StartInfo.WorkingDirectory = directory.FullName;
			p.StartInfo.UseShellExecute = false;
			p.StartInfo.RedirectStandardOutput = true;
			p.StartInfo.RedirectStandardError = true;

			p.Start();
			p.WaitForExit();

			if (p.ExitCode != 0)
			{
				throw new Exception(string.Format("Running phonopy failed: {0}", p.StandardError.ReadToEnd()));
			}

			List<Structure> displacedStructures = new List<Structure>();
			foreach (FileInfo file in directory.GetFiles())
			{
				if (file.Name.StartsWith("POSCAR-"))
				{
					displacedStructures.Add(VaspPoscar.Read(file));
				}
			}

			return new PhonopyProblem(baseStructure, repeatMatrix, displacement, VaspPoscar.Read(directory.AddFile("SPOSCAR")), displacedStructures, CompressedData.Read(directory.AddFile("disp.yaml")), CompressedData.Read(directory.AddFile("POSCAR")));
		}

		public static string FormatRepeatMatrix(LatticeRepeatMatrix repeatMatrix)
		{
			int m11 = repeatMatrix[0, 0];
			int m12 = repeatMatrix[0, 1];
			int m13 = repeatMatrix[0, 2];
			int m21 = repeatMatrix[1, 0];
			int m22 = repeatMatrix[1, 1];
			int m23 = repeatMatrix[1, 2];
			int m31 = repeatMatrix[2, 0];
			int m32 = repeatMatrix[2, 1];
			int m33 = repeatMatrix[2, 2];

			if (m12 == 0 && m13 == 0 && m21 == 0 && m23 == 0 && m31 == 0 && m32 == 0)
			{
				return string.Format("{0} {1} {2}", m11, m22, m33);
			}
			else
			{
				return string.Format("{0} {1} {2}  {3} {4} {5}  {6} {7} {8}", m11, m12, m13, m21, m22, m23, m31, m32, m33);
			}
		}

		/// <summary>
		/// Extracts the computed forces and create a matrix of force constants. Must be in the same order as the displaced structures
		/// in DisplacedStructures.
		/// </summary>
		public PhonopyResult Collect(params VaspResult[] results)
		{
			return Collect((IEnumerable<VaspResult>)results);
		}

		/// <summary>
		/// Extracts the computed forces and create a matrix of force constants. Must be in the same order as the displaced structures
		/// in DisplacedStructures.
		/// </summary>
		public PhonopyResult Collect(IEnumerable<VaspResult> results)
		{
			using (TempDirectory temp = new TempDirectory())
			{
				return Collect(results, temp.Directory);
			}
		}

		public PhonopyResult CollectResidual(VaspResult result0, params VaspResult[] results)
		{
			return CollectResidual(result0, (IEnumerable<VaspResult>)results);
		}

		public PhonopyResult CollectResidual(VaspResult result0, IEnumerable<VaspResult> results)
		{
			using (TempDirectory temp = new TempDirectory())
			{
				return CollectResidual(result0, results, temp.Directory);
			}
		}

		private PhonopyResult Collect(IEnumerable<VaspResult> results, DirectoryInfo directory)
		{
			List<string> fileNames = new List<string>();
			int n = 0;

			foreach (VaspResult result in results)
			{
				DirectoryInfo d = directory.AddDirectory("disp-" + n.ToString().PadLeft(3, '0'));
				d.Create();
				File.WriteAllBytes(d.AddFile("vasprun.xml").FullName, result.Vasprun.GetData());

				fileNames.Add(d.Name + "/" + "vasprun.xml");
			}

			File.WriteAllBytes(directory.AddFile("disp.yaml").FullName, Disp.GetData());

			Process p = new Process();
			p.StartInfo.FileName = "phonopy";
			p.StartInfo.Arguments = "-f " + string.Join(" ", fileNames);
			p.StartInfo.WorkingDirectory = directory.FullName;
			p.StartInfo.UseShellExecute = false;
			p.StartInfo.RedirectStandardOutput = true;
			p.StartInfo.RedirectStandardError = true;

			p.Start();
			p.WaitForExit();

			if (p.ExitCode != 0)
			{
				throw new Exception(string.Format("Running phonopy failed: {0}", p.StandardError.ReadToEnd()));
			}

			return new PhonopyResult(this, CompressedData.Read(directory.AddFile("FORCE_SETS")));
		}

		private PhonopyResult CollectResidual(VaspResult result0, IEnumerable<VaspResult> results, DirectoryInfo directory)
		{
			List<string> fileNames = new List<string>();
			int n = 0;

			DirectoryInfo d0 = directory.AddDirectory("sposcar");
			d0.Create();
			File.WriteAllBytes(d0.AddFile("vasprun.xml").FullName, result0.Vasprun.GetData());
			fileNames.Add(d0.Name + "/" + "vasprun.xml");

			foreach (VaspResult result in results)
			{
				DirectoryInfo d = directory.AddDirectory("disp-" + n.ToString().PadLeft(3, '0'));
				d.Create();
				File.WriteAllBytes(d.AddFile("vasprun.xml").FullName, result.Vasprun.GetData());

				fileNames.Add(d.Name + "/" + "vasprun.xml");
			}

			File.WriteAllBytes(directory.AddFile("disp.yaml").FullName, Disp.GetData());

			Process p = new Process();
			p.StartInfo.FileName = "phonopy";
			p.StartInfo.Arguments = "--fz " + string.Join(" ", fileNames);
			p.StartInfo.WorkingDirectory = directory.FullName;
			p.StartInfo.UseShellExecute = false;
			p.StartInfo.RedirectStandardOutput = true;
			p.StartInfo.RedirectStandardError = true;

			p.Start();
			p.WaitForExit();

			if (p.ExitCode != 0)
			{
				throw new Exception(string.Format("Running phonopy failed: {0}", p.StandardError.ReadToEnd()));
			}

			return new PhonopyResult(this, CompressedData.Read(directory.AddFile("FORCE_SETS")));
		}

		public Structure BaseStructure
		{
			get;
			private set;
		}

		public LatticeRepeatMatrix RepeatMatrix
		{
			get;
			private set;
		}

		public double? Displacement
		{
			get;
			private set;
		}

		/// <summary>
		/// The super structure on which the finite displacements are performed. No computation is needed on this.
		/// </summary>
		public Structure SuperStructure
		{
			get;
			private set;
		}

		/// <summary>
		/// The finite displacement structures to be computed in order to determine the force constants.
		/// </summary>
		public ReadOnlyList<Structure> DisplacedStructures
		{
			get;
			private set;
		}

		/// <summary>
		/// Phonopy's disp.yaml file.
		/// </summary>
		public CompressedData Disp
		{
			get;
			private set;
		}

		/// <summary>
		/// Phonopy's POSCAR file.
		/// </summary>
		public CompressedData Poscar
		{
			get;
			private set;
		}
	}
}
